Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[XLA:GPU] Support partially pipelined async send recv ops #17446

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

copybara-service[bot]
Copy link

[XLA:GPU] Support partially pipelined async send recv ops

This is needed for pipeline parallelism on GPU where the send/recv operations
are issued in one loop iteration and completed in the next. The same buffer
must be alive throughout the process and no copies can be inserted.

Avoid copies for these partially pipelined async send/recv ops. Insert the
required copies and controlflow constraints on the send/recv ops separately.
This is to ensure that the live times of the buffers do not overlap.

Send: For send, a copy is inserted on the operand, starting a new live range.
By enforcing this copy after the corresponding send/done, buffer live times are
disjoint.

Recv: For recv, a copy is inserted after recv-done, ending the live time of the
buffer. Bt enforcing the copy to be before the corresponding recv. buffer live
times are disjoint.

This is needed for pipeline parallelism on GPU where the send/recv operations
are issued in one loop iteration and completed in the next. The same buffer
must be alive throughout the process and no copies can be inserted.

Avoid copies for these partially pipelined async send/recv ops. Insert the
required copies and controlflow constraints on the send/recv ops separately.
This is to ensure that the live times of the buffers do not overlap.

Send: For send, a copy is inserted on the operand, starting a new live range.
By enforcing this copy after the corresponding send/done, buffer live times are
disjoint.

Recv: For recv, a copy is inserted after recv-done, ending the live time of the
buffer. Bt enforcing the copy to be before the corresponding recv. buffer live
times are disjoint.
PiperOrigin-RevId: 676075106
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant